package com.llmops.core;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

public class WorkflowGraph extends Node {
    private static final Logger logger = LoggerFactory.getLogger(WorkflowGraph.class);
    private String parentId;
    private final String id;
    private final Map<String, Node> nodes = new ConcurrentHashMap<>();
    private final List<Edge> edges = new CopyOnWriteArrayList<>();
    private Object result;
    public WorkflowGraph(String id) {
        super(id, null);
        this.id = id;
    }

    public WorkflowGraph(String parentId, String id) {
        super(id, null);
        this.id = id;
        this.parentId = parentId;
    }

    public void addNode(Node node) {
        nodes.put(node.getId(), node);
    }

    public void addEdge(String sourceId, String targetId, Edge.ConditionType type, Object value) {
        Edge edge = new Edge(sourceId, targetId, type, value);
        edges.add(edge);
        Node targetNode = nodes.get(targetId);
        if (targetNode != null) {
            targetNode.previousNodes.add(nodes.get(sourceId));
        }
    }

    public void addEdge(Edge edge) {
        edges.add(edge);
        Node targetNode = nodes.get(edge.getTargetId());
        if (targetNode != null) {
            targetNode.previousNodes.add(nodes.get(edge.getSourceId()));
        }
    }

    public void validate() throws IllegalStateException {
        if (nodes.isEmpty()) {
            throw new IllegalStateException("Empty workflow graph");
        }
        // 拓扑排序验证逻辑预留
    }

    @Override
    public void execute(Context context) {
        Queue<Node> executionQueue = new LinkedList<>();
        //获取起始节点
        nodes.values().stream()
                .filter(node -> edges.stream().noneMatch(e -> e.getTargetId().equals(node.getId())))
                .forEach(executionQueue::add);

        while (!executionQueue.isEmpty()) {
            Worker worker = executionQueue.poll();
            Node current = (Node) worker;
            current.setStatus(Node.NodeStatus.RUNNING);
            context.setNodeStatus(current.getId(), Node.NodeStatus.RUNNING);
            current.preExecute(context);
            try {
                logger.info("开始=====>[type:{}][id:{}]", current.getType(), current.getId());
                current.execute(context);
                logger.info("结束=====>[type:{}][id:{}]", current.getType(), current.getId());
                current.setStatus(Node.NodeStatus.SUCCESS);
                context.setNodeStatus(current.getId(), Node.NodeStatus.SUCCESS);
                context.setVariable("result",context.getVariable(current.getId()+"_output"));
            } catch (Exception e) {
                current.setStatus(Node.NodeStatus.ERROR);
                context.setNodeStatus(current.getId(), Node.NodeStatus.ERROR);
                throw e;
            } finally {
                current.postExecute(context);
            }

            edges.stream()
                    .filter(e -> e.getSourceId().equals(current.getId()))
                    .filter(e -> e.evaluateCondition(context))
                    .map(e -> nodes.get(e.getTargetId()))
                    .filter(Objects::nonNull)
                    .forEach(executionQueue::add);
        }
    }

//    public Context getContext(){
//        return context;
//    }


    public String getParentId() {
        return parentId;
    }

    @Override
    public String getType() {
        return "workFlow";
    }

    public String getId() {
        return id;
    }

    public Map<String, Node> getNodes() {
        return nodes;
    }

    public List<Edge> getEdges() {
        return edges;
    }
}